"""
11_rag_agent.py - RAG Agent with LangGraph
============================================
Chapter 13: Building intelligent retrieval-augmented generation.

KEY CONCEPTS:
  ★ Query Generation - Expand user question into search queries
  ★ Retrieval Loop - Retry until confidence threshold met
  ★ Context Assembly - Combine retrieved docs for generation
  ★ Confidence Check - Evaluate retrieval quality

Run: python src/11_rag_agent.py
"""
from typing import TypedDict, Annotated
from operator import add
from langgraph.graph import StateGraph, START, END

class RAGState(TypedDict):
    question: str
    search_queries: list
    retrieved_docs: Annotated[list, add]
    answer: str
    confidence: float
    iteration: int

def generate_queries(state: RAGState) -> dict:
    q = state["question"]
    return {"search_queries": [q, f"{q} explained", f"{q} examples"]}

def retrieve_documents(state: RAGState) -> dict:
    docs = []
    for query in state["search_queries"]:
        docs.append({"query": query, "content": f"Information about: {query}", "score": 0.85})
    return {"retrieved_docs": docs}

def evaluate_docs(state: RAGState) -> dict:
    scores = [d["score"] for d in state["retrieved_docs"]]
    avg = sum(scores) / max(len(scores), 1)
    return {"confidence": avg, "iteration": state.get("iteration", 0) + 1}

def generate_answer(state: RAGState) -> dict:
    context = "\n".join(d["content"] for d in state["retrieved_docs"][:3])
    return {"answer": f"Based on {len(state['retrieved_docs'])} sources:\n{context}"}

def should_retry(state: RAGState) -> str:
    if state.get("confidence", 0) >= 0.8 or state.get("iteration", 0) >= 3:
        return "generate"
    return "retry"

graph = StateGraph(RAGState)
graph.add_node("gen_queries", generate_queries)
graph.add_node("retrieve", retrieve_documents)
graph.add_node("evaluate", evaluate_docs)
graph.add_node("generate", generate_answer)
graph.add_edge(START, "gen_queries")
graph.add_edge("gen_queries", "retrieve")
graph.add_edge("retrieve", "evaluate")
graph.add_conditional_edges("evaluate", should_retry, {"generate": "generate", "retry": "gen_queries"})
graph.add_edge("generate", END)
app = graph.compile()

if __name__ == "__main__":
    print("=" * 60)
    print("11: RAG Agent")
    print("=" * 60)
    result = app.invoke({"question": "What is LangGraph?"})
    print(f"\n  Answer: {result['answer']}")
    print(f"  Confidence: {result['confidence']:.2f}")
    print(f"  Iterations: {result['iteration']}")
    print(f"  Docs retrieved: {len(result['retrieved_docs'])}")
    print(f"\n{'=' * 60}\n✅ RAG agent complete!")
